from bolt_nut_detection_model_trainer import BoltNutDetectionModelTrainer
import argparse
import os
from six.moves import urllib
import zipfile
import ssl
import re

_DIR_PATH = ""
_PRETRAINED_MODEL = "pretrained-yolov3.h5"

def parse_args():
    parser = argparse.ArgumentParser(description='Train YOLO networks with random input shape.')
    parser.add_argument('--batch-size', type=int, default=4, help='Training mini-batch size')
    parser.add_argument('--dataset', type=str, default='voc', help='Training dataset. Now support voc.')
    parser.add_argument('--pre-model', dest="pre_model", type=str, help='Pretrained model.')
    parser.add_argument('--data-url', dest='data_url', type=str, default='', help='Training dataset download url.')
    parser.add_argument('--epochs', type=int, default=100, help='Training epochs.')
    args = parser.parse_args()

    return args

def download_train_pretrained_model(data_url):
    print('pre_model:', data_url)
    filepath = os.path.join(_DIR_PATH, _PRETRAINED_MODEL)
    filepath, _ = urllib.request.urlretrieve(data_url, filepath)

def download_dataset(data_url):
    print('data_url:', data_url)
    filepath = os.path.join(_DIR_PATH)
    filepath, _ = urllib.request.urlretrieve(data_url, filepath)

    with zipfile.ZipFile(filepath, 'r') as zip_file:
        zip_file.extractall(_DIR_PATH)

if __name__ == '__main__':
    args = parse_args()

    download_train_pretrained_model(args.pre_model)
    download_dataset(args.dataset)

    from keras import backend as K

    # set GPU memory
    if ('tensorflow' == K.backend()):
        import tensorflow as tf
        from keras.backend.tensorflow_backend import set_session

        config = tf.ConfigProto()
        config.gpu_options.per_process_gpu_memory_fraction = 0.8
        config.gpu_options.allow_growth = True
        sess = tf.Session(config=config)

    trainer = BoltNutDetectionModelTrainer()
    trainer.setModelTypeAsYOLOv3()
    trainer.setDataDirectory(data_directory="boltNut")
    trainer.setTrainConfig(object_names_array=["螺母", "螺栓"], batch_size=args.batch_size, num_experiments=200,
                           train_from_pretrained_model=_PRETRAINED_MODEL)
    trainer.trainModel()


